#include "bi_traversal.hpp"

namespace hms::swe
{

template<typename MeshType>
void computeCells(
	SolverThread<MeshType>& threadData,
	const IndexHelper& ih
){
	using B = BlockLocation;
	#define COMPUTE_CELL(LOC, BEG_X, BEG_Y) \
	threadData.template computeBlock<true>( Block<B::LOC>{ih, BEG_X, BEG_Y, 1, 1} )

	Index cellBegX, cellBegY;
	const Index
		& nxInner { ih.mesh.nxInner },
		& nyInner { ih.mesh.nyInner };

	/* inner cells */
	OMP_PRAGMA( for schedule(static) collapse(2) nowait )
	for (Index j=0; j<nyInner; ++j ){
	for (Index i=0; i<nxInner; ++i ){
		cellBegX = ih.beg.blocks + i;
		cellBegY = ih.beg.blocks + j;
		COMPUTE_CELL(inner, cellBegX, cellBegY);
	}}

	/* bottom and top rows, without corners */
	#define COMPUTE_ROW(LOC,Y_BEG) \
	OMP_PRAGMA( for schedule(dynamic) nowait ) \
	for ( Index i=0; i<nxInner; ++i )\
		COMPUTE_CELL(LOC, ih.beg.blocks + i, Y_BEG);

	/* left and right cols, without corners */
	#define COMPUTE_COL(LOC,X_BEG) \
	OMP_PRAGMA( for schedule(dynamic) nowait ) \
	for ( Index i=0; i<nyInner; ++i )\
		COMPUTE_CELL(LOC, X_BEG, ih.beg.blocks + i);

	/* outer bottom */
	COMPUTE_ROW(bottom, ih.beg.mesh);
	/* inner bottom */
	COMPUTE_ROW(bottom, ih.beg.mesh + 1);
	/* inner top */
	COMPUTE_ROW(top, ih.beg.outerY);
	/* outer bottom */
	COMPUTE_ROW(top, ih.beg.outerY + 1);

	/* outer left */
	COMPUTE_COL(left, ih.beg.mesh);
	/* inner left */
	COMPUTE_COL(left, ih.beg.mesh + 1);
	/* inner right */
	COMPUTE_COL(right, ih.beg.outerX);
	/* outer right */
	COMPUTE_COL(right, ih.beg.outerX + 1);

	#undef COMPUTE_ROW
	#undef COMPUTE_COL

	/* remainder blocks and corners */
	/* removed nowait clause, because it caused issues with the swap of old and
	 * new fields. And generally, after all blocks have been computed, a
	 * barrier does make sense. */
	OMP_PRAGMA( sections ) // nowait
	{
		#define COMPUTE_CORNER(LOC, X_BEG, Y_BEG) \
		OMP_PRAGMA(section) \
		{ \
			COMPUTE_CELL(LOC, X_BEG  , Y_BEG  ); \
			COMPUTE_CELL(LOC, X_BEG+1, Y_BEG  ); \
			COMPUTE_CELL(LOC, X_BEG  , Y_BEG+1); \
			COMPUTE_CELL(LOC, X_BEG+1, Y_BEG+1); \
		}

		COMPUTE_CORNER(bottomLeft , ih.beg.mesh  , ih.beg.mesh  )
		COMPUTE_CORNER(bottomRight, ih.beg.outerX, ih.beg.mesh  )
		COMPUTE_CORNER(topLeft    , ih.beg.mesh  , ih.beg.outerY)
		COMPUTE_CORNER(topRight   , ih.beg.outerX, ih.beg.outerY)
		
		#undef COMPUTE_CORNER
	}

	#undef COMPUTE_CELL
}

template void computeCells( SolverThread< UniMesh>&, const IndexHelper& );
template void computeCells( SolverThread<RectMesh>&, const IndexHelper& );

} // namespace hms::swe